2017/Day 07.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dataclasses import dataclass, field\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class Program(object):\n",
    "    weight: int\n",
    "    dependents: set = field(default_factory=set)\n",
    "    _total_weight: int = field(init=False, default=None)\n",
    "\n",
    "    def get_total_weight(self, programs):\n",
    "        if self._total_weight is None:\n",
    "            self._total_weight = self.weight + sum(\n",
    "                programs[d].get_total_weight(programs) for d in self.dependents\n",
    "            )\n",
    "        return self._total_weight\n",
    "\n",
    "\n",
    "def read_programs(lines):\n",
    "    programs = {}\n",
    "    for line in lines:\n",
    "        if not line.strip():\n",
    "            continue\n",
    "        name_weight, arrow, dependents = line.partition(\" -> \")\n",
    "        name, weight = name_weight.split(\" (\")\n",
    "        programs[name] = Program(\n",
    "            int(weight.rstrip(\")\\n \")),\n",
    "            {d.strip() for d in dependents.split(\",\") if d.strip()},\n",
    "        )\n",
    "    return programs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import aocd\n",
    "\n",
    "data = aocd.get_data(day=7, year=2017)\n",
    "programs = read_programs(data.splitlines())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def find_root(program):\n",
    "    roots = program.keys() - set.union(*(v.dependents for v in program.values()))\n",
    "    return roots.pop()\n",
    "\n",
    "\n",
    "test = read_programs(\n",
    "    \"\"\"\\\n",
    "pbga (66)\n",
    "xhth (57)\n",
    "ebii (61)\n",
    "havc (66)\n",
    "ktlj (57)\n",
    "fwft (72) -> ktlj, cntj, xhth\n",
    "qoyq (66)\n",
    "padx (45) -> pbga, havc, qoyq\n",
    "tknk (41) -> ugml, padx, fwft\n",
    "jptl (61)\n",
    "ugml (68) -> gyxo, ebii, jptl\n",
    "gyxo (61)\n",
    "cntj (57)\n",
    "\"\"\".splitlines()\n",
    ")\n",
    "assert find_root(test) == \"tknk\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "tests = {\"ugml\": 251, \"padx\": 243, \"fwft\": 243}\n",
    "\n",
    "for name, expected in tests.items():\n",
    "    assert test[name].get_total_weight(test) == expected\n",
    "\n",
    "\n",
    "def correct_balance(name, programs):\n",
    "    program = programs[name]\n",
    "    weights = {}\n",
    "    for d in program.dependents:\n",
    "        weight = programs[d].get_total_weight(programs)\n",
    "        weights.setdefault(weight, []).append(d)\n",
    "    if len(weights) == 1:\n",
    "        # balanced, no adjustment needed\n",
    "        return 0\n",
    "\n",
    "    # imbalanced\n",
    "    imbalanced, target = weights\n",
    "    if len(weights[imbalanced]) != 1:\n",
    "        imbalanced, target = target, imbalanced\n",
    "\n",
    "    # Check for balance in the child nodes\n",
    "    sub_program_name = weights[imbalanced][0]\n",
    "    sub_correction = correct_balance(sub_program_name, programs)\n",
    "    if sub_correction:\n",
    "        return sub_correction\n",
    "\n",
    "    return programs[sub_program_name].weight + (target - imbalanced)\n",
    "\n",
    "\n",
    "assert correct_balance(\"tknk\", test) == 60"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Part 1: gynfwly\n"
     ]
    }
   ],
   "source": [
    "root = find_root(programs)\n",
    "print(\"Part 1:\", root)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Part 2: 1526\n"
     ]
    }
   ],
   "source": [
    "print(\"Part 2:\", correct_balance(root, programs))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}